package textprocessing.nlpanalysis.constituent;

import algorithm.config.HumanLanguage;
import algorithm.nlp.corenlp.TextAnnotator;
import edu.stanford.nlp.ling.CoreAnnotations;
import edu.stanford.nlp.ling.Label;
import edu.stanford.nlp.pipeline.Annotation;
import edu.stanford.nlp.trees.Tree;
import edu.stanford.nlp.trees.TreeCoreAnnotations;
import edu.stanford.nlp.util.CoreMap;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

/**
 * Created by common on 2017/4/8.
 */
public class SCITreeGenerator {

    static final Logger logger = LoggerFactory.getLogger(SCITreeGenerator.class);

    private SCITree.SCITreeNode lastNPNode;
    private SCITree.SCITreeNode leftNPNode;
    private Tree leftMostNPNode;

    public SCITreeGenerator(){

    }



    public List<SCITree> getSCITreeOfDocument(Annotation document){

        List<CoreMap> sentences = document.get(CoreAnnotations.SentencesAnnotation.class);
        List<SCITree> sciTreeList = new ArrayList<>();

        for(CoreMap sentence: sentences) {
            List<SCITree> senList = getSCITreeOfSentence(sentence);
            sciTreeList.addAll(senList);
        }
        return sciTreeList;
    }


    public List<SCITree> getSCITreeOfSentence(CoreMap sentence){
        // 初始化全局变量
        this.lastNPNode = null;
        this.leftMostNPNode = null;

        Tree tree = sentence.get(TreeCoreAnnotations.TreeAnnotation.class);
        tree.pennPrint();

        SCITree sciTree = new SCITree();
        SCITree.SCITreeNode root = sciTree.root;

        /*
        最左边的NP
         */
        leftmost_NP(tree);

        //traversal_recursive( tree, root);
        traversal_marksub(tree, root);
        System.out.println("\n\n\n\n-------------------" + "Original SCI Tree After Replace Label" + "-----------------------");
        SCITree.traversal_print( root , "");
        //System.out.println("\n\n\n\n");


        List<SCITree> subSciTreeList = new ArrayList<>();
        subSciTreeList.add(sciTree);


        // 抽取出子句对应的子树
        extract_subclause(root, subSciTreeList);


        // 对于每棵子树
        for(SCITree sub : subSciTreeList){
            // 标记 ENUM CONC
            traversal_subtree( sub.root);
            // 合并叶子节点
            merge_concatenation( sub.root );
        }

        return subSciTreeList;
    }


    /*
    标记子树 SUB
     */
    private void traversal_marksub(Tree tree, SCITree.SCITreeNode root){
        // 标记原始节点
        root.originalNode = tree;

        if(tree == this.leftMostNPNode){
            this.leftNPNode = root;
        }

        Label label = tree.label();
        String labelRaw = label.value();

        if(!tree.isLeaf()) {

            // 保存深度优先遍历过程中，最近出现的NP，作为从句的主语
            if (labelRaw.equalsIgnoreCase("NP")) {
                this.lastNPNode = root;
            }

            /*
            （SCR 2)
             */
            if(labelRaw.equalsIgnoreCase("S") || labelRaw.equalsIgnoreCase("SBARQ ")){
                root.nodeType = SCITree.NodeType.SUB;
            }

            /*
            (SCI 3): Mark as SUB each SBAR
             */
            if (labelRaw.equalsIgnoreCase("SBAR")) {
                root.nodeType = SCITree.NodeType.SUB;
                // add a head NP(Noun Phrase) as Subject, search backward.
                if (this.lastNPNode != null) {
                    root.headNode = this.lastNPNode;
                    System.out.println("\n\n[Find SBAR] Last NP Node = " + this.lastNPNode.originalNode + "\n\n");
                }
            }



            /*
            (SCI 4): Mark as SUB each PP starting with a preposition
            EG: he loves apple but hates apple pie
             */
            // TODO: check does this reasonable?
            if (labelRaw.equalsIgnoreCase("PP")) {
                String firstLabel = tree.firstChild().label().value();
                // TODO: ONLY sonme kinds of PP should become context
                if(firstLabel.equalsIgnoreCase("CC")){
                    root.nodeType = SCITree.NodeType.SUB;
                    // add subject of sentence to headNode
                    // get left-most NP in the Sub-Clause
                    root.headNode = this.leftNPNode;
                }
            }

        }

        else{ //叶子节点
            root.nodeType = SCITree.NodeType.LEAF;
            root.isLeaf = true;

            root.content = labelRaw;

            // 创建叶子节点的 context
            root.contexts.add( labelRaw );
            ContextItem item = new ContextItem();
            item.addTokenFirst(root.originalNode);
            root.contextItems.add(item);

        }

        /*
        DFS -> 深度优先遍历
         */
        for(Tree child : tree.children() ){
            String childLabel = child.label().value();
            /*
            TODO： ADD RULE
             */
            if(childLabel.equalsIgnoreCase("CC")){  //连词不在结果中输出
                continue;
            }
            SCITree.SCITreeNode childNode = new SCITree.SCITreeNode();
            root.childs.add(childNode);
            traversal_marksub(child, childNode);
        }
    }



    /*
    分割子树
     */
    private void extract_subclause(SCITree.SCITreeNode root, List<SCITree> subList){
        if(root.nodeType == SCITree.NodeType.SUB){
            SCITree sciTree = new SCITree();
            sciTree.root = root;
            subList.add(sciTree);
        }

        for(SCITree.SCITreeNode sub: root.childs){
            extract_subclause(sub, subList);
        }

        /*
         如何在迭代过程中删除元素。
         http://www.jianshu.com/p/c5b52927a61a
         */
        Iterator<SCITree.SCITreeNode> iter = root.childs.iterator();
        while(iter.hasNext()){
            SCITree.SCITreeNode sub = iter.next();
            if( sub.nodeType == SCITree.NodeType.SUB )
            {
                iter.remove();
            }
        }

        /*
        for(SCITree.SCITreeNode sub: root.childs){
            if(sub.nodeType == SCITree.NodeType.SUB){
                root.childs.remove( sub );
            }
        }
        */
    }


    /*
    标记每棵子树  ENUM CONC
     */
    private void traversal_subtree( SCITree.SCITreeNode root){

        Tree tree = root.originalNode;
        Label label = root.originalNode.label();
        String labelRaw = label.value();

        System.out.println("[Current Node] ------ node = " + tree);

        if(!tree.isLeaf()) {

            // 保存深度优先遍历过程中，最近出现的NP，作为从句的主语
            if (labelRaw.equalsIgnoreCase("NP")) {
                this.lastNPNode = root;
            }

            /*
            **********************************************
            * 注意：规则的应用是有序的，不能随意打乱顺序
            * 另外，假设子树已经拆分完毕
             */

            /*
            (SCI 5): Mark as CONC remaining nodes
             */
            root.nodeType = SCITree.NodeType.CONC;

            /*
            (SCI 1): unless children are [all NP] or [all VP]
             */
            boolean all_np = true;
            boolean all_vp = true;
            int valid_child = 0;
            boolean contain_con_or_pun = false;
            for (Tree child : tree.children()) {
                String childLabel = child.label().value();
                System.out.println("[Check All Child] child = " + child + " \nchildLabel = " + childLabel);
                if (childLabel.equalsIgnoreCase("CC")) {  //连接词，例如and,or,等，忽略
                    contain_con_or_pun = true;
                    continue;
                }
                // TODO: 检测标点符号
                else if (childLabel.equalsIgnoreCase(".")
                        || childLabel.equalsIgnoreCase(",")
                        || childLabel.equalsIgnoreCase("?")
                        ) { //punctuation
                    contain_con_or_pun = true;
                    continue;
                }
                else {
                    if (!childLabel.equalsIgnoreCase("NP")) {
                        all_np = false;
                    }
                    if (!childLabel.equalsIgnoreCase("VP")) {
                        all_vp = false;

                    }
                    if(childLabel.equalsIgnoreCase("NP") || childLabel.equalsIgnoreCase("VP"))
                        valid_child += 1;
                }
            }
            if ((all_np || all_vp) && valid_child>0 && contain_con_or_pun) {
                root.nodeType = SCITree.NodeType.ENUM;
                System.out.println("[NodeType = ENUM]\n");
            }


        }

        else{ //叶子节点, 已经生成过
            /*
            root.nodeType = SCITree.NodeType.LEAF;
            root.isLeaf = true;

            root.content = labelRaw;

            // 创建叶子节点的 context
            root.contexts.add( labelRaw );
            ContextItem item = new ContextItem();
            item.addTokenFirst(root.originalNode);
            root.contextItems.add(item);
            */
        }


        /*
        DFS -> 深度优先遍历
         */
        for(SCITree.SCITreeNode childNode : root.childs ){
            //String childLabel = child.label().value();
            /*
            TODO： ADD RULE
             */
            /*
            if(childLabel.equalsIgnoreCase("CC")){  //连词不在结果中输出
                continue;
            }
            */
            //SCITree.SCITreeNode childNode = new SCITree.SCITreeNode();
            //root.childs.add(childNode);

            traversal_subtree(childNode);
        }
    }




    /*
    该函数已经被拆分了
     */
    private void traversal_recursive(Tree tree, SCITree.SCITreeNode root){
        // 标记原始节点
        root.originalNode = tree;

        if(tree == this.leftMostNPNode){
            this.leftNPNode = root;
        }

        Label label = tree.label();
        String labelRaw = label.value();

        if(!tree.isLeaf()) {

            // 保存深度优先遍历过程中，最近出现的NP，作为从句的主语
            if (labelRaw.equalsIgnoreCase("NP")) {
                this.lastNPNode = root;
            }


            /*
            **********************************************
            * 注意：规则的应用是有序的，不能随意打乱顺序
             */


            /*
            (SCI 4): Mark as CONC remaining nodes
             */
            root.nodeType = SCITree.NodeType.CONC;

            /*
            (SCI 1): unless children are [all NP] or [all VP]
             */
            boolean all_np = true;
            boolean all_vp = true;
            int valid_child = 0;
            for (Tree child : tree.children()) {
                String childLabel = child.label().value();
                if (childLabel.equalsIgnoreCase("CC")) {  //连接词，例如and,or,等，忽略
                    continue;
                }
                // TODO: 检测标点符号
                else if (childLabel.equalsIgnoreCase(".") || childLabel.equalsIgnoreCase(",")) { //punctuation
                    continue;
                } else {
                    if (!childLabel.equalsIgnoreCase("NP")) {
                        all_np = false;
                    }
                    else if (!childLabel.equalsIgnoreCase("VP")) {
                        all_vp = false;
                    }
                    if(childLabel.equalsIgnoreCase("NP") || childLabel.equalsIgnoreCase("VP"))
                        valid_child += 1;
                }
            }
            if ((all_np || all_vp) && valid_child>0) {
                root.nodeType = SCITree.NodeType.ENUM;
                //System.out.println("[NodeType = ENUM]\n\n\n\n");
            }

            /*
            (SCI 2): Mark as SUB each SBAR
             */
            if (labelRaw.equalsIgnoreCase("SBAR")) {
                root.nodeType = SCITree.NodeType.SUB;
                // add a head NP(Noun Phrase) as Subject, search backward.
                if (this.lastNPNode != null) {
                    root.headNode = this.lastNPNode;
                }
            }

            if(labelRaw.equalsIgnoreCase("S") || labelRaw.equalsIgnoreCase("SBARQ ")){
                root.nodeType = SCITree.NodeType.SUB;
            }


            /*
            (SCI 3): Mark as SUB each PP starting with a preposition
            EG: he loves apple but hates apple pie
             */
            // TODO: check does this reasonable?
            if (labelRaw.equalsIgnoreCase("PP")) {
                String firstLabel = tree.firstChild().label().value();
                // TODO: ONLY sonme kinds of PP should become context
                if(firstLabel.equalsIgnoreCase("CC")){
                    root.nodeType = SCITree.NodeType.SUB;
                    // add subject of sentence to headNode
                    // get left-most NP in the Sub-Clause
                    root.headNode = this.leftNPNode;
                }
            }

        }

        else{ //叶子节点
            root.nodeType = SCITree.NodeType.LEAF;
            root.isLeaf = true;

            root.content = labelRaw;

            // 创建叶子节点的 context
            root.contexts.add( labelRaw );
            ContextItem item = new ContextItem();
            item.addTokenFirst(root.originalNode);
            root.contextItems.add(item);
        }


        /*
        DFS -> 深度优先遍历
         */
        for(Tree child : tree.children() ){
            String childLabel = child.label().value();
            /*
            TODO： ADD RULE
             */
            if(childLabel.equalsIgnoreCase("CC")){  //连词不在结果中输出
                continue;
            }
            SCITree.SCITreeNode childNode = new SCITree.SCITreeNode();
            root.childs.add(childNode);

            traversal_recursive(child, childNode);
        }
    }


    private Tree leftmost_NP(Tree tree){
        this.leftMostNPNode = null;
        Tree leftmost_np_node = null;
        do{
            String label = tree.label().value();
            if(label.equalsIgnoreCase("NP"))
                leftmost_np_node = tree;
            if(!tree.isLeaf()){
                tree = tree.firstChild();
            }
        }
        while(!tree.isLeaf());
        if(leftmost_np_node == null){
            leftmost_np_node = tree;
        }
        this.leftMostNPNode = leftmost_np_node;
        return leftmost_np_node;
    }







    /*
    把所有子节点也是叶子节点的中间节点，处理为叶子节点
     */
    private boolean merge_concatenation(SCITree.SCITreeNode root){
        boolean is_leaf = false;

        if(root.isLeaf){
            is_leaf = true;
        }
        else{
            // 后序遍历，从叶子节点往上合并
            boolean all_leaf_child = true;
            String content = "";
            ContextItem leaf_item = new ContextItem();
            for(SCITree.SCITreeNode childNode : root.childs){
                try{
                    if(childNode.nodeType == SCITree.NodeType.SUB ){
                        continue;
                        // SUB会被踢出去，成为单独一棵树
                    }
                    boolean leaf_child = merge_concatenation(childNode);
                    if(!leaf_child)
                        all_leaf_child = false;
                    else{
                        content += (" " + childNode.content);
                        if(childNode.contextItems.size() != 1){
                            logger.error("[LeafNode] Context Num Error! Node Info = \n{}", childNode.originalNode.toString());
                            logger.error("size = {}, label = {}", childNode.contextItems.size(), childNode.originalNode.label());
                            throw new Exception("Context Num Error");
                        }
                        leaf_item.extendLast( childNode.contextItems.get(0) );
                    }
                }catch (Exception e){
                    e.printStackTrace();
                }

            }

            if(all_leaf_child && root.nodeType== SCITree.NodeType.CONC){
                //标记当前节点为叶子节点
                root.content = content;
                root.isLeaf = true;
                is_leaf = true;
                // 构造 context
                root.contextItems.add( leaf_item );

            }

        }

        return is_leaf;
    }



    /*
    *************************************************************
    * UNIT TEST
     */
    public static void main(String[] args){
        String text = "Barack Obama was a lawyer before he became the president.";

        TextAnnotator textAnnotator = new TextAnnotator(HumanLanguage.ENGLISH);
        Annotation document = textAnnotator.getAnnotatedText(text);


        SCITreeGenerator generator = new SCITreeGenerator();
        generator.getSCITreeOfDocument(document);
    }


}
